#pragma once

#include "compi/collectives.h"
#include "compi/context.h"
#include "compi/p2p.h"
#include "compi/result.h"
#include "compi/serializer.h"
#include "compi/tag.h"

#include <optional>

#include <atomic>
#include <memory>
#include <mpi.h>
#include <mutex>
#include <string>
#include <unordered_map>

namespace compi
{
// Endpoint：仅持有通信器与目标 rank，便于替换后端
class Endpoint
{
public:
    Endpoint(MPI_Comm comm, int target_rank) noexcept : _comm(comm), _rank(target_rank) {}

    MPI_Comm comm() const noexcept { return _comm; }
    int rank() const noexcept { return _rank; }

private:
    MPI_Comm _comm;
    int _rank;
};

class FramedChannel
{
public:
    explicit FramedChannel(bool allow_any_tag = true) noexcept : _allow_any_tag(allow_any_tag) {}

    void set_failure_policy_override(std::optional<FailurePolicy> policy) noexcept
    {
        _failure_override = policy;
    }

    std::optional<FailurePolicy> failure_policy_override() const noexcept
    {
        return _failure_override;
    }

    Result<bool> send_frame(const Endpoint & dst,
                            const std::vector<std::byte> & frame_bytes,
                            int tag = 0) noexcept;

    Result<bool> recv_frame(const Endpoint & src,
                            std::vector<std::byte> & out_frame_bytes,
                            int tag = 0) noexcept;

    Result<bool> send_frame(const Endpoint & dst,
                            const std::vector<std::byte> & frame_bytes,
                            TagNamespace ns) noexcept;

    Result<bool> recv_frame(const Endpoint & src,
                            std::vector<std::byte> & out_frame_bytes,
                            TagNamespace ns) noexcept;

    TransportRequest isend_frame(const Endpoint & dst,
                                 const std::vector<std::byte> & frame_bytes,
                                 TagNamespace ns,
                                 SendMode mode = SendMode::Default) noexcept;

    TransportRequest isend_frame(const Endpoint & dst,
                                 const std::vector<std::byte> & frame_bytes,
                                 const TagLease & lease,
                                 SendMode mode = SendMode::Default) noexcept;

    TransportRequest irecv_frame(const Endpoint & src,
                                 std::vector<std::byte> & out_frame_bytes,
                                 TagNamespace ns) noexcept;

    bool allow_any_tag() const noexcept { return _allow_any_tag; }
    void set_allow_any_tag(bool v) noexcept { _allow_any_tag = v; }

private:
    bool _allow_any_tag{true};
    std::optional<FailurePolicy> _failure_override{};
};

template <typename T>
class TypedChannel
{
public:
    explicit TypedChannel(FramedChannel & framed) noexcept : _framed(framed) {}

    Result<bool> send(const Endpoint & dst, const T & payload, int tag = 0) noexcept
    {
        try
        {
            auto bytes = serialize_to_byte_vector<T>(payload);
            return _framed.send_frame(dst, bytes, tag);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

    Result<bool> recv(const Endpoint & src, T & payload, int tag = 0) noexcept
    {
        std::vector<std::byte> bytes;
        auto res = _framed.recv_frame(src, bytes, tag);
        if (res.is_error())
            return res;
        try
        {
            payload = deserialize<T>(bytes);
            return Result<bool>::success(true);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

    Result<bool> send(const Endpoint & dst, const T & payload, TagNamespace ns) noexcept
    {
        try
        {
            auto bytes = serialize_to_byte_vector<T>(payload);
            return _framed.send_frame(dst, bytes, ns);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

    Result<bool> recv(const Endpoint & src, T & payload, TagNamespace ns) noexcept
    {
        std::vector<std::byte> bytes;
        auto res = _framed.recv_frame(src, bytes, ns);
        if (res.is_error())
            return res;
        try
        {
            payload = deserialize<T>(bytes);
            return Result<bool>::success(true);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

private:
    FramedChannel & _framed;
};

} // namespace compi

namespace compi
{

// 设计目标（Channel 层）：
// - 基于 Context 的统一 MPI 调用入口（可恢复错误栈）
// - 显式 TLV 帧协议集成（Framing）
// - 每 communicator 真实单例的 TagAllocator（资源管理）
// - 可选 per-channel communicator（隔离或共享通信域）

struct ChannelId
{
    std::string name;
    std::size_t hash_value;

    explicit ChannelId(const std::string & channel_name)
        : name(channel_name), hash_value(std::hash<std::string>{}(channel_name))
    {
    }

    bool operator==(const ChannelId & other) const noexcept
    {
        return hash_value == other.hash_value && name == other.name;
    }

    std::size_t hash() const noexcept { return hash_value; }
};

// 集中通道选项：用于统一配置 Channel 层行为
struct ChannelOptions
{
    // 是否为每个通道复制独立 communicator（MPI_Comm_dup）。
    // 若为 false，则复用上层 Context/Manager 的 communicator。
    bool use_dedicated_comm = true;

    // 默认命名空间名称（可选）。当通过命名空间发送/接收帧时，未显式指定时使用此名称。
    // 由 MessageTagAllocator 在通道激活后解析与管理；空字符串表示不设置默认命名空间。
    std::string default_namespace_name{};

    // ANY_TAG 行为：是否允许使用 ANY_TAG 进行接收（如 recv_frame_any）。
    // 为 false 时，应限制为命名空间内的标签或显式 tag。
    bool allow_any_tag = true;

    // 失败策略覆盖：是否覆盖 Context 的 FailurePolicy。
    bool override_failure_policy = false;
    FailurePolicy failure_policy = FailurePolicy::ThrowException;

};

enum class ChannelState
{
    Created,
    Active,
    Closing,
    Closed,
};

// 统计信息（面向通道的轻量指标）
struct ChannelStats
{
    std::size_t messages_sent = 0;
    std::size_t messages_received = 0;
    std::size_t bytes_sent = 0;
    std::size_t bytes_received = 0;
    double total_time_ms = 0.0;
};

// 通道：围绕 Context 与 Framing 的统一消息通道
// - 可选择独立通信器（MPI_Comm_dup）或复用 Context 的 comm
// - 所有 MPI 操作通过 Context::call_mpi，错误进入 ErrorStack，可按策略抛出/继续
// - 标签通过 MessageTagAllocator 单例管理（按 communicator）
class Channel
{
public:
    // 工厂：创建通道（可选独立通信器）
    static Result<std::shared_ptr<Channel>>
    create(const ChannelId & id, MPI_Comm parent_comm, bool use_dedicated_comm = true) noexcept;

    ~Channel();

    Channel(const Channel &) = delete;
    Channel & operator=(const Channel &) = delete;
    Channel(Channel &&) = delete;
    Channel & operator=(Channel &&) = delete;

    const ChannelId & id() const noexcept { return _id; }
    MPI_Comm comm() const noexcept { return _comm; }
    Context & context() noexcept { return *_context; }

    // 选项：集中配置的访问与更新
    const ChannelOptions & options() const noexcept { return _options; }
    void set_options(const ChannelOptions & opts) noexcept
    {
        const std::string previous_default_ns = _options.default_namespace_name;
        _options = opts;
        _framed.set_allow_any_tag(_options.allow_any_tag);
        if (_options.override_failure_policy)
        {
            _framed.set_failure_policy_override(_options.failure_policy);
        }
        else
        {
            _framed.set_failure_policy_override(std::nullopt);
        }
        if (previous_default_ns != _options.default_namespace_name)
        {
            // default namespace will be (re)initialized on next activate
            _default_namespace.reset();
            _owns_default_namespace = false;
            if (is_active())
            {
                auto res = ensure_default_namespace();
                if (res.is_error() && _context)
                {
                    _context->handle_error(res);
                }
            }
        }
    }

    int rank() const noexcept { return _context ? _context->rank() : -1; }
    int size() const noexcept { return _context ? _context->size() : -1; }

    ChannelState state() const noexcept { return _state.load(); }

    // 激活/关闭（幂等）
    Result<void> activate() noexcept;
    Result<void> close() noexcept;

    bool is_active() const noexcept { return _state.load() == ChannelState::Active; }
    bool is_closed() const noexcept { return _state.load() == ChannelState::Closed; }

    // 统计
    const ChannelStats & stats() const noexcept { return _stats; }
    void reset_stats() noexcept { _stats = {}; }

    // 资源管理：按 communicator 的真实单例分配器
    MessageTagAllocator & tag_allocator() noexcept;
    std::optional<TagNamespace> default_namespace_handle() const noexcept { return _default_namespace; }

    // Framing：发送/接收统一帧字节流（不再使用旧 tlv::Frame）
    Result<bool>
    send_frame(int dest_rank, const std::vector<std::byte> & frame_bytes, int tag = 0) noexcept;
    Result<bool>
    recv_frame(int src_rank, std::vector<std::byte> & out_frame_bytes, int tag = 0) noexcept;

    // Framing（ANY_TAG 过滤 + 命名空间）：可选基于命名空间管理消息标签（TagNamespace）
    Result<bool>
    send_frame(int dest_rank, const std::vector<std::byte> & frame_bytes, TagNamespace ns) noexcept;
    Result<bool>
    recv_frame(int src_rank, std::vector<std::byte> & out_frame_bytes, TagNamespace ns) noexcept;

    TransportRequest isend_frame_async(int dest_rank,
                                       const std::vector<std::byte> & frame_bytes,
                                       TagNamespace ns,
                                       SendMode mode = SendMode::Default) noexcept;
    TransportRequest isend_frame_async(int dest_rank,
                                       const std::vector<std::byte> & frame_bytes,
                                       const TagLease & lease,
                                       SendMode mode = SendMode::Default) noexcept;
    TransportRequest irecv_frame_async(int src_rank,
                                       std::vector<std::byte> & out_frame_bytes,
                                       TagNamespace ns) noexcept;

    // Typed IO：类型安全的发送/接收（序列化到字节帧），委托给 FramedChannel
    template <typename T>
    Result<bool> send(int dest_rank, const T & payload, int tag = 0) noexcept
    {
        try
        {
            auto bytes = serialize_to_byte_vector<T>(payload);
            if (auto * ns = default_namespace(); ns && tag == 0)
            {
                return send_frame(dest_rank, bytes, *ns);
            }
            Endpoint ep(_comm, dest_rank);
            auto res = _framed.send_frame(ep, bytes, tag);
            if (res.is_error())
                return res;
            _stats.messages_sent++;
            _stats.bytes_sent += bytes.size();
            return Result<bool>::success(true);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

    template <typename T>
    Result<bool> recv(int src_rank, T & payload, int tag = 0) noexcept
    {
        Endpoint ep(_comm, src_rank);
        std::vector<std::byte> bytes;
        Result<bool> recv_result = Result<bool>::success(true);
        if (auto * ns = default_namespace(); ns && tag == 0)
        {
            recv_result = _framed.recv_frame(ep, bytes, *ns);
        }
        else
        {
            recv_result = _framed.recv_frame(ep, bytes, tag);
        }
        if (recv_result.is_error())
            return recv_result;
        try
        {
            payload = deserialize<T>(bytes);
            _stats.messages_received++;
            _stats.bytes_received += bytes.size();
            return Result<bool>::success(true);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

    template <typename T>
    Result<bool> send(int dest_rank, const T & payload, TagNamespace ns) noexcept
    {
        try
        {
            auto bytes = serialize_to_byte_vector<T>(payload);
            Endpoint ep(_comm, dest_rank);
            auto res = _framed.send_frame(ep, bytes, ns);
            if (res.is_error())
                return res;
            _stats.messages_sent++;
            _stats.bytes_sent += bytes.size();
            return Result<bool>::success(true);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

    template <typename T>
    Result<bool> recv(int src_rank, T & payload, TagNamespace ns) noexcept
    {
        Endpoint ep(_comm, src_rank);
        std::vector<std::byte> bytes;
        auto res = _framed.recv_frame(ep, bytes, ns);
        if (res.is_error())
            return res;
        try
        {
            payload = deserialize<T>(bytes);
            _stats.messages_received++;
            _stats.bytes_received += bytes.size();
            return Result<bool>::success(true);
        }
        catch (const std::exception & e)
        {
            return Result<bool>::error(ResultCode::SerializationError, e.what());
        }
    }

private:
    Channel(const ChannelId & id, MPI_Comm comm, Context * ctx, bool owns_comm) noexcept;
    Result<void> initialize_comm_properties() noexcept;
    Result<void> ensure_default_namespace() noexcept;
    Result<void> release_default_namespace() noexcept;
    TagNamespace * default_namespace() noexcept
    {
        return _default_namespace ? &*_default_namespace : nullptr;
    }
    TransportRequest track_request(TransportRequest req) noexcept;
    void prune_completed_requests_locked() noexcept;

private:
    ChannelId _id;
    MPI_Comm _comm;
    Context * _context;
    bool _owns_comm;
    std::atomic<ChannelState> _state{ChannelState::Created};
    ChannelStats _stats{};
    ChannelOptions _options{}; // 集中通道选项（默认值）
    FramedChannel _framed{};   // 帧层委托（承载 ANY_TAG 策略）
    std::optional<TagNamespace> _default_namespace{};
    bool _owns_default_namespace{false};
    mutable std::mutex _pending_mutex;
    std::vector<TransportRequest> _pending_requests;
};

// 管理器：每个 MPI_Comm 一个实例，负责通道的创建/查询/销毁
class ChannelManager
{
public:
    static ChannelManager & for_comm(MPI_Comm comm);
    static void destroy_for_comm(MPI_Comm comm);

    Result<std::shared_ptr<Channel>> create_channel(const ChannelId & id,
                                                    bool use_dedicated_comm = true);
    // 以集中选项创建通道：在创建后立即写入选项
    Result<std::shared_ptr<Channel>> create_channel(const ChannelId & id,
                                                    const ChannelOptions & opts);
    std::shared_ptr<Channel> get_channel(const ChannelId & id) const;
    Result<void> destroy_channel(const ChannelId & id);

    std::size_t active_channel_count() const;
    void cleanup_all();

    // 将析构函数设为 public，以便共享持有（shared_ptr）能够安全删除
    ~ChannelManager();

private:
    explicit ChannelManager(MPI_Comm comm);

    ChannelManager(const ChannelManager &) = delete;
    ChannelManager & operator=(const ChannelManager &) = delete;

    struct ChannelIdHash
    {
        std::size_t operator()(const ChannelId & id) const noexcept { return id.hash(); }
    };

    MPI_Comm _comm;
    std::unordered_map<ChannelId, std::shared_ptr<Channel>, ChannelIdHash> _channels;
    mutable std::mutex _mutex;
};

// 组：批量激活/关闭与资源清理（围绕同一 Context / communicator）
class ChannelGroup
{
public:
    explicit ChannelGroup(const std::string & group_name);
    ~ChannelGroup();

    ChannelGroup(const ChannelGroup &) = delete;
    ChannelGroup & operator=(const ChannelGroup &) = delete;

    bool add_channel(std::shared_ptr<Channel> channel) noexcept;
    bool remove_channel(const ChannelId & id) noexcept;
    std::shared_ptr<Channel> get_channel(const ChannelId & id) const noexcept;

    std::size_t activate_all() noexcept;
    std::size_t close_all() noexcept;

    std::size_t size() const noexcept;
    const std::string & name() const noexcept { return _name; }

private:
    std::string _name;
    struct ChannelIdHash
    {
        std::size_t operator()(const ChannelId & id) const noexcept { return id.hash(); }
    };
    std::unordered_map<ChannelId, std::shared_ptr<Channel>, ChannelIdHash> _channels;
    mutable std::mutex _mutex;
};

} // namespace compi
